import paras
from torch.autograd import Variable
from madin import *
import matplotlib.pyplot as plt
import numpy as np
import random
import autoencoder as ae

span = 100
predspan = 0
# torchfilename = 'data/FXUSDJPY_autoencoder.pkl'
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# device = 'cpu'


def processInput(trandata):
    tranmax = trandata.max(1)
    tranmin = trandata.min(1)
    transpan = tranmax - tranmin
    trandata = (trandata.T - tranmin)
    return (trandata / transpan).T * 2 - 1


def autocode(dbname="FX", symbol="FXUSDJPY", qt_type=5):
    # fetchdata.fetchdata(symbol=symbol, qt_type=qt_type)
    # datalist = fetchdata.loaddate(fetchdata.abspath + "data/%s.json" % symbol)
    numbers=paras.number
    datalist = finddb(dbname=dbname, tablename=symbol, limit=numbers)
    # datalist = json.load(open("data/A%s.json" % 'FXUSDJPY', 'r'))
    # start=-int(numbers/qt_type)
    end=int(numbers/qt_type*0.75)
    datalist = joinKline(datalist, qt_type)[0:end]
    # datalist = finddb(dbname="stock", tablename=symbol, limit=20000)
    # torchfilename = 'data/%s_autoencoder.pkl' % symbol if qt_type == 5 else 'data/%s_autoencoder_1.pkl' % symbol
    torchfilename = 'data/{}_{}_autoencoder.pkl'.format(symbol, qt_type)
    # datalist=datalist[-2200:]
    trandata = []
    trandata2 = []
    for i in range(len(datalist) - span - predspan):
        line = []
        for j in datalist[i + 1:i + 1 + span]:
            line.append(j["Close"])
        trandata.append(line)
    for i in range(len(datalist) - span - predspan):
        line = []
        for j in datalist[i + 1:i + 1 + span + predspan]:
            line.append(j["Close"])
        trandata2.append(line)

    weight_u = pow(3, 1 / span)
    weightarray = np.array([weight_u ** i * 5 if i > span else weight_u ** i for i in range(0, span + predspan)])
    weightarray = weightarray / np.sum(weightarray)
    # weightarray.reshape(-1,1)
    # weightarray=np.column_stack([weightarray*0.5,weightarray*0.2,weightarray*0.2,weightarray*2])
    # weightarray=weightarray.reshape(-1,1)
    weightarray = Variable(torch.from_numpy(weightarray).view(-1, 1).type(torch.FloatTensor))

    showpos = [random.randint(0, len(trandata) - 1) for i in range(9)]
    trandata = np.array(trandata, dtype=np.float32)
    trandata = processInput(trandata)
    trandata2 = np.array(trandata2, dtype=np.float32)
    trandata2 = processInput(trandata2)

    try:
        autoencoder = torch.load(torchfilename).to(device)
    except:
        autoencoder = ae.AutoEncoder(span).to(device)
    autoencoder.train()
    tran1 = torch.from_numpy(trandata).to(device)
    tran2 = torch.from_numpy(trandata2).to(device)
    # tran2=tran1

    f, axs = plt.subplots(3, 3)
    axlist = []
    for li in axs:
        for o in li:
            axlist.append(o)

    optimizer = torch.optim.Adam(autoencoder.parameters())  # 传入 net 的所有参数, 学习率
    autoencoder.train()
    loop = 0
    lossmin = 1e10
    while loop < 20:
        autoencoder.train()
        encoded, decoded = autoencoder(tran1)  # 喂给 net 训练数据 x, 输出预测值
        # print(decoded.shape,tran2.shape)

        # loss=torch.mean(torch.mm(torch.pow(decoded-tran2,2),weightarray))
        loss = torch.mean(torch.pow(decoded - tran2, 2))

        optimizer.zero_grad()  # 清空上一步的残余更新参数值
        loss.backward()  # 误差反向传播, 计算参数更新值
        optimizer.step()  # 将参数更新值施加到 net 的 parameters 上
        print('loop={},Loss={}'.format(loop, loss.item()))

        if lossmin > loss.item():
            lossmin = loss.item()
            loop = 0
        else:
            loop += 1

        if loop % 10 == 0:
            autoencoder.eval()
            encoded, decoded = autoencoder(tran1)
            d1, d2 = decoded.data.cpu().numpy(), tran2.data.cpu().numpy()
            xd = np.arange(d1.shape[1])
            for o in range(len(showpos)):
                axlist[o].cla()
                axlist[o].plot(xd, d1[showpos[o]], 'r', lw=1)
                axlist[o].plot(xd, d2[showpos[o]], 'g', lw=1)
            plt.pause(0.01)

    plt.close()
    autoencoder.eval()
    torch.save(autoencoder.cpu(), torchfilename)


if __name__ == "__main__":
    autocode(paras.dbname, paras.symbol, paras.qt_type)
